"""
使用rerank模型来实现document的rerank的实现比较简单
"""
from util.T2llm import T2llm


def run(
        query: str,
        documents: list[str],
        score_threshold: float = None,
        top_n: int = None
) -> list[tuple]: #(document,score) //文档和分数 排序之后的

    # 调用rerank模型对文档进行重排，返回每个文档的分数
    llm = T2llm()
    scores = llm.rerank(query, documents)

    rerank_documents = []
    # 过滤掉分数最低的
    for document, score in zip(documents, scores):
        if score_threshold and score < score_threshold:
            continue
        rerank_documents.append((document, score))

    # 将所有符合要求的文档按分数从高到低排序。
    rerank_documents.sort(key=lambda x: x[1], reverse=True)

    # 根据 top_n 参数返回评分最高的前 N 个文档；如果不设置 top_n，则返回所有符合条件的文档。
    return rerank_documents[:top_n] if top_n else rerank_documents

if __name__ == '__main__':
    query = "苹果"
    documents = ["苹果","大苹果","香蕉","梨子"]
    result = run(query,documents,top_n=3)
    print(result)
    # [('苹果', tensor(7.7574)), ('大苹果', tensor(3.1842)), ('梨子', tensor(-0.5926))]
